package com.ustc.hewei.nettyrpc2_0.zookeeper.register;

import com.ustc.hewei.nettyrpc2_0.common.warp.ZkNode;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.x.discovery.ServiceDiscovery;
import org.apache.curator.x.discovery.ServiceInstance;
import org.apache.curator.x.discovery.ServiceInstanceBuilder;

import java.net.InetAddress;
import java.util.concurrent.ConcurrentHashMap;

/**
* @description: zookeeper注册类
* @author hewei
* @date 2022/9/15 20:05
* @version 1.0
*/

@Slf4j
public class ZkServiceRegister extends DefaultServiceRegister implements ServiceRegister{
    private final ServiceDiscovery<ZkNode> serviceDiscovery;
    private final ConcurrentHashMap<String, ServiceInstance<ZkNode>> instanceMap;

    public ZkServiceRegister(int port, ServiceDiscovery<ZkNode> serviceDiscovery) {
        this.port = port;
        this.serviceDiscovery = serviceDiscovery;
        instanceMap = new ConcurrentHashMap<>();
    }

    @Override
    public void register(Object obj, String serviceName, int weight) {
        try {
            super.register(obj, serviceName, weight);
            ZkNode zkNode = new ZkNode();
            String hostname = InetAddress.getLocalHost().getHostAddress();
            zkNode.setHost(hostname).setCur(0)
                    .setPort(port).setServiceName(serviceName).setWeight(weight);
            ServiceInstanceBuilder<ZkNode> builder = ServiceInstance.builder();
            ServiceInstance<ZkNode> serviceInstance = builder.payload(zkNode)
                    .name(serviceName).port(zkNode.getPort())
                    .address(zkNode.getHost()).build();
            serviceDiscovery.registerService(serviceInstance);
            instanceMap.put(serviceName, serviceInstance);
            log.info("服务 {} 注册成功", zkNode);
        } catch (Exception e) {
            throw new RuntimeException("注册服务时异常", e);
        }
    }

    @Override
    public void unregister(String serviceName) throws Exception {
        super.unregister(serviceName);
        serviceDiscovery.unregisterService(instanceMap.remove(serviceName));
    }
}